import pandas as pd
import os
import argparse

def is_one_to_one(row):
    for item in row:
        # If item is NaN or empty string, return False
        if not isinstance(item, str) or item == '':
            return False
        # If there are multiple genes for this species, return False
        if len(item.split(',')) > 1:
            return False
    return True


def main(args):
    # 1. 加载所有必要的数据
    zscore_folder = args.zscore_folder
    orthogroups_file = args.orthogroups_file
    meta_data_file = args.meta_data_file
    outfile = args.outfile

    # Load OGS data
    ogs_df = pd.read_csv(orthogroups_file, sep="\t", index_col='Orthogroup')

    # Load meta data
    meta_df = pd.read_csv(meta_data_file, sep="\t")
    print(meta_df)

    # 2. 建立需要的数据结构
    label_mapping = {}
    for index, row in meta_df.iterrows():
        label_mapping[row["sample_name"]] = f"{row['label']}@{row['spe']}@{row['sample_name']}"

    # 3. 查找1:1:1基因
    one_to_one_genes = ogs_df[ogs_df.iloc[:, :].apply(is_one_to_one, axis=1)]

    # 4. 创建最终的数据矩阵
    result_df = pd.DataFrame()
    result_df["OGS_name"] = one_to_one_genes.index

    for file in os.listdir(zscore_folder):
        if file.endswith(".tsv") or file.endswith(".csv"):
            species = file.split("_")[0]
            zscore_df = pd.read_csv(os.path.join(zscore_folder, file), sep="\t")
            for col in zscore_df.columns[1:]:
                if col in label_mapping:
                    formatted_col_name = label_mapping[col]
                    #print(formatted_col_name)
                    # Reverse the mapping
                    gene_to_orthogroup_mapping = one_to_one_genes[species].reset_index().set_index(species)['Orthogroup'].to_dict()
                    # Now map the gene_id to its Orthogroup
                    zscore_df["OGS"] = zscore_df["gene_id"].map(gene_to_orthogroup_mapping)
                    subset = zscore_df.dropna(subset=["OGS"])[["OGS", col]]
                    #print(zscore_df["OGS"])
                    subset.set_index("OGS", inplace=True)
                    
                    result_df[formatted_col_name] = result_df["OGS_name"].map(subset[col])

    result_df = result_df.dropna()
    result_df.to_csv(outfile, sep="\t", index=False)



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process zscore tables and create a final matrix for SVM.")
    parser.add_argument("zscore_folder", type=str, help="Directory containing zscore tables.")
    parser.add_argument("orthogroups_file", type=str, help="Path to the Orthogroups_gene_name.tsv file.")
    parser.add_argument("meta_data_file", type=str, help="Path to the meta_data.tsv file.")
    parser.add_argument("outfile", type=str, help="Path to the output file (final_matrix_for_SVM.tsv).")

    args = parser.parse_args()
    main(args)